{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**The code for runs 1 to 4**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from chainer.datasets import mnist\n",
    "import brian2 as b\n",
    "from brian2tools import *\n",
    "import matplotlib.pyplot as plt\n",
    "from brian2 import NeuronGroup, Subgroup, Synapses\n",
    "import matplotlib.cm as cmap\n",
    "import os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "b.prefs.codegen.target= 'cython' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "testMode=False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#split data in train and test\n",
    "train, test = mnist.get_mnist(withlabel=True, ndim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#show one image\n",
    "x, t = train[0]\n",
    "plt.imshow(x.reshape(28, 28), cmap='gray')\n",
    "plt.savefig('5.png')\n",
    "print('label:', t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#create PoissonGroup to start with\n",
    "startGroup = b.PoissonGroup((28*28), 0*b.Hz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Number of neurons used to create the two neuron groups in the processing layer\n",
    "#Number of excitatory and inhibitory neurons should be equal\n",
    "\n",
    "nInput = 28*28 #number of input neurons, 784\n",
    "nE = 400 #number of excitatory neurons, can be changed\n",
    "nI = nE #number of inhibitory neurons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#These equations define the models of the two neuron groups\n",
    "\n",
    "neuronEquationsE = '''\n",
    "        dv/dt = ((restE - v) + (I_synE+I_synI) / nS) / (100*second)  : volt\n",
    "        I_synE = ge * nS *         -v                           : amp\n",
    "        I_synI = gi * nS * (-100.*mV-v)                         : amp\n",
    "        dge/dt = -ge/(1.0*second)                               : 1\n",
    "        dgi/dt = -gi/(2.0*second)                               : 1\n",
    "        '''\n",
    "if(testMode==False):\n",
    "    neuronEquationsE += '\\n  dtheta/dt = -theta / (tcTheta)  : volt' \n",
    "else:\n",
    "    neuronEquationsE += '\\n theta : volt'\n",
    "\n",
    "neuronEquationsE += '\\n  dtimer/dt = 100.0  : second'\n",
    "\n",
    "neuronEquationsI = '''\n",
    "        dv/dt = ((restI - v) + (I_synE+I_synI) / nS) / (10*second)  : volt\n",
    "        I_synE = ge * nS *         -v                               : amp\n",
    "        I_synI = gi * nS * (-85.*mV-v)                              : amp\n",
    "        dge/dt = -ge/(1.0*second)                                   : 1\n",
    "        dgi/dt = -gi/(2.0*second)                                   : 1\n",
    "        '''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create neuron groups\n",
    "#Set parameters \n",
    "\n",
    "neuronGroups = {}\n",
    "\n",
    "refracE = 5. * b.ms\n",
    "refracI = 2. * b.ms\n",
    "resetE = -65. * b.mV\n",
    "resetI = -45. * b.mV\n",
    "threshE = -52. * b.mV\n",
    "threshI = -40. * b.mV\n",
    "\n",
    "tcTheta = 1e7 * b.ms\n",
    "thetaPlusE = 0.05 * b.mV\n",
    "offset = 20.0*b.mV\n",
    "\n",
    "vResetI = 'v=-45.*mV' \n",
    "\n",
    "#theta represents adaptive threshold --> intrinsic plasticity\n",
    "vResetE = 'v = resetE; theta += thetaPlusE; timer = 0*ms' \n",
    "\n",
    "if(testMode==True):\n",
    "    vResetE = 'v = resetE; timer = 0*ms'\n",
    "\n",
    "vThreshI = 'v>-40.*mV'\n",
    "\n",
    "vThreshE = '(v>(theta - offset + -52.*mV)) and (timer>refracE)'\n",
    "\n",
    "#Make excitatory and inhibitory neurons in the processing layer\n",
    "neuronGroups['e'] = b.NeuronGroup(nE, neuronEquationsE, threshold= vThreshE, refractory= refracE, reset= vResetE) \n",
    "neuronGroups['i'] = b.NeuronGroup(nI, neuronEquationsI, threshold= vThreshI, refractory= refracI, reset= vResetI)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Resting membrane potentials for the neuron subgroups\n",
    "\n",
    "restE = -65. * b.mV \n",
    "restI = -60. * b.mV "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create neuron subgroups\n",
    "\n",
    "neuronGroups['Pe'] = Subgroup(source=neuronGroups['e'], start=0, stop=nE)\n",
    "neuronGroups['Pi'] = Subgroup(source= neuronGroups['i'], start=0, stop=nI)\n",
    "neuronGroups['Pe'].v = restE - 40. * b.mV \n",
    "neuronGroups['Pi'].v = restI - 40. * b.mV"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set the thetha for the neuron groups\n",
    "\n",
    "weight_path='./'\n",
    "\n",
    "if(testMode==False):\n",
    "    neuronGroups['Pe'].theta = np.ones((nE)) * 20.0*b.mV #only excitatory\n",
    "#else:\n",
    "#Load own saved thetas, check brian1 from their code\n",
    "#we did not implement this, since we didn't get further than just training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create lateral connections in processing layer\n",
    "#Use random weights to start with\n",
    "\n",
    "connections = {}\n",
    "\n",
    "name = 'PePi'\n",
    "connections[name] = Synapses(neuronGroups['Pe'], neuronGroups['Pi'],'w: 1')\n",
    "connections[name].connect(i=neuronGroups['Pe'], j=neuronGroups['Pi'])\n",
    "connections[name].w='rand()'\n",
    "\n",
    "\n",
    "connName = 'PiPe'\n",
    "connections[name] = Synapses(neuronGroups['Pi'], neuronGroups['Pe'],'w: 1')\n",
    "connections[name].connect(i=neuronGroups['Pi'], j=neuronGroups['Pe']) \n",
    "connections[name].w='rand()'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create the monitors\n",
    "#Rate monitors for Pe and Pi\n",
    "#Spike monitor only for Pe\n",
    "\n",
    "#Spikes for Pi are not connected directly to Xe, so are not relevant to measure\n",
    "\n",
    "rateMonitors = {}\n",
    "spikeMonitors = {}\n",
    "\n",
    "rateMonitors['Pe'] = b.PopulationRateMonitor(neuronGroups['Pe'])\n",
    "rateMonitors['Pi'] = b.PopulationRateMonitor(neuronGroups['Pi'])\n",
    "spikeMonitors['Pe'] = b.SpikeMonitor(neuronGroups['Pe'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create input population and add rate monitor\n",
    "\n",
    "inputGroups = {}\n",
    "inputGroups['Xe'] = b.PoissonGroup(nInput, 0*b.Hz)\n",
    "rateMonitors['Xe'] = b.PopulationRateMonitor(inputGroups['Xe'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create input connections\n",
    "\n",
    "tc_pre_ee = 20*b.ms \n",
    "tc_post_1_ee = 20*b.ms\n",
    "tc_post_2_ee = 40*b.ms \n",
    "nu_ee_pre =  0.0001      # learning rate \n",
    "nu_ee_post = 0.01       # learning rate \n",
    "wmax_ee = 1.0 \n",
    "\n",
    "eqs_stdp_pre_ee = 'pre = 1.; w -= nu_ee_pre * post1'\n",
    "eqs_stdp_post_ee = 'post2before = post2; w += nu_ee_post * pre * post2before; post1 = 1.; post2 = 1.'\n",
    "\n",
    "wmax_ee = 1.0 #weights cannot be higher than 1\n",
    "\n",
    "#Add STDP only for training\n",
    "if(testMode==False):\n",
    "    connections['XePe'] = Synapses(inputGroups['Xe'], neuronGroups['Pe'], \n",
    "                '''w:1\n",
    "                    post2before                            : 1\n",
    "                    dpre/dt   =   -pre/(tc_pre_ee)         : 1 (event-driven)\n",
    "                    dpost1/dt  = -post1/(tc_post_1_ee)     : 1 (event-driven)\n",
    "                    dpost2/dt  = -post2/(tc_post_2_ee)     : 1 (event-driven)\n",
    "                ''',\n",
    "                on_pre= eqs_stdp_pre_ee,\n",
    "                on_post= eqs_stdp_post_ee) #triplet STDP\n",
    "else:\n",
    "    connections['XePe'] = Synapses(inputGroups['Xe'], neuronGroups['Pe'])\n",
    "\n",
    "connections['XePe'].connect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set intervals for the simulation!\n",
    "\n",
    "oneExampleTime =   0.35 * b.second #each Poisson spike train is presented 350ms\n",
    "restingTime = 0.15 * b.second #150ms resting phase before new image is presented, to allow neurons to decay to their resting values\n",
    "\n",
    "nExamples = 500 * 1\n",
    "\n",
    "if(testMode==True):\n",
    "    nExamples = 10000*1\n",
    "    updateInterval = nExamples\n",
    "\n",
    "runtime = nExamples * (oneExampleTime + restingTime)\n",
    "if nExamples <= 10000:    \n",
    "    updateInterval = nExamples\n",
    "    weightUpdateInterval = 20\n",
    "else:\n",
    "    updateInterval = 10000\n",
    "    weightUpdateInterval = 100\n",
    "if nExamples <= 60000:    \n",
    "    saveConnectionsInterval = 10000\n",
    "else:\n",
    "    saveConnectionsInterval = 10000\n",
    "    updateInterval = 10000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set input intensity to 2\n",
    "#Is increased later on until spike count threshold is reached\n",
    "#Once this is reached, it is set back to startInputIntensity\n",
    "\n",
    "inputIntensity = 2.\n",
    "startInputIntensity = inputIntensity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plotPerformanceFigure(fig):\n",
    "    nEvaluations = int(nExamples/updateInterval) #determine how often to plot performance\n",
    "    steps = range(0, nEvaluations) #determine the time steps\n",
    "    performance = np.zeros(nEvaluations) #start empty for every step\n",
    "    figure = b.figure(fig, figsize = (5, 5))\n",
    "    fig += 1\n",
    "    ax = figure.add_subplot(111)\n",
    "    im2, = ax.plot(steps, performance) \n",
    "    b.ylim(ymax = 100)\n",
    "    b.title('Classification performance')\n",
    "    figure.canvas.draw()\n",
    "    return im2, performance, fig, figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#the weight normalization from Goodhill and Barrow\n",
    "\n",
    "def normalizeWeights():\n",
    "    connection = np.zeros((nInput, nE))\n",
    "    connection[connections['XePe'].i, connections['XePe'].j] = connections['XePe'].w\n",
    "    temp = np.copy(connection)\n",
    "    colSums = np.sum(temp, axis = 0) #sum of all weights\n",
    "    colFactors = 78./colSums #weight of the ee_input?\n",
    "    for j in range(nE):\n",
    "        temp[:,j] *= colFactors[j] #multiply all weights by the factor\n",
    "    connections['XePe'].w = temp[connections['XePe'].i, connections['XePe'].j] #update the weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train, test = mnist.get_mnist(withlabel=True, ndim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "#The simulation\n",
    "\n",
    "net = b.Network()\n",
    "for obj_list in [neuronGroups, inputGroups, connections, spikeMonitors]:\n",
    "    for key in obj_list:\n",
    "        net.add(obj_list[key])\n",
    "\n",
    "previousSpikeCount = np.zeros(nE)\n",
    "inputNumbers = [0] * nExamples\n",
    "inputGroups['Xe'].rates = 0 * b.Hz\n",
    "net.run(0*b.second)\n",
    "j = 0\n",
    "while(j < (int(nExamples))):\n",
    "    if(testMode==False):\n",
    "        normalizeWeights()\n",
    "        rate = train[j%60000][0].reshape((nInput)) / 8. *  inputIntensity\n",
    "    else:\n",
    "        rate = train[j%10000][0].reshape((nInput)) / 8. *  inputIntensity\n",
    "    \n",
    "    inputGroups['Xe'].rates = rate * b.kHz\n",
    "    net.run(oneExampleTime)\n",
    "\n",
    "    currentSpikeCount = np.asarray(spikeMonitors['Pe'].count[:]) - previousSpikeCount\n",
    "    previousSpikeCount = np.copy(spikeMonitors['Pe'].count[:])\n",
    "    \n",
    "    #Increase intensity if there are less than 5 spikes\n",
    "    if(np.sum(currentSpikeCount) < 5):\n",
    "        inputIntensity += 1\n",
    "        inputGroups['Xe'].rates = 0 * b.Hz\n",
    "        net.run(restingTime)\n",
    "    else:\n",
    "        if(j % 100 == 0 and j > 0):\n",
    "            print('runs done:', j, 'of', int(nExamples))\n",
    "        inputGroups['Xe'].rates = 0 * b.Hz\n",
    "        net.run(restingTime)\n",
    "        inputIntensity = startInputIntensity\n",
    "        j += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Save results\n",
    "#Not actually tested since we never got this far\n",
    "\n",
    "if(testMode==False):\n",
    "    os.makedirs(os.path.dirname(weight_path+'ourWeightsTest4200/thetaP'), exist_ok=True)\n",
    "    os.makedirs(os.path.dirname(weight_path+'ourWeightsTest4200/XePe'), exist_ok=True)\n",
    "    os.makedirs(os.path.dirname(weight_path+'ourWeightsTest4200/PePi'), exist_ok=True)\n",
    "    os.makedirs(os.path.dirname(weight_path+'ourWeightsTest4200/PiPe'), exist_ok=True)\n",
    "    \n",
    "    np.save(weight_path+'ourWeightsTest4200/thetaP', neuronGroups['Pe'].theta)\n",
    "    connectionList = zip(connections['XePe'].i, connections['XePe'].j, connections['XePe'].w)\n",
    "    connectionList2 = zip(connections['PePi'].i, connections['PePi'].j, connections['PePi'].w)\n",
    "    connectionList3 = zip(connections['PiPe'].i, connections['PiPe'].j, connections['PiPe'].w)\n",
    "\n",
    "    np.save(weight_path+'ourWeightsTest4200/XePe', np.array(list(zip(connectionList)))\n",
    "    np.save(weight_path+'ourWeightsTest4200/PePi', np.array(list(zip(connectionList2)))\n",
    "    np.save(weight_path+'ourWeightsTest4200/PiPe', np.array(list(zip(connectionList3)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
